import os
from sentence_transformers import SentenceTransformer, InputExample, losses, models
from torch.utils.data import DataLoader
from sentence_transformers.evaluation import TripletEvaluator
from sentence_transformers.util import cos_sim
import random

# --- 步骤 1: 准备训练数据 ---
# 在真实场景中，这些数据应该从你的文件中加载
# 这里我们用一个简化的样本数据集
print("--- 1. 准备训练数据 ---")
data = {
    "九寨沟": [
        "九寨沟以其多彩的湖泊、瀑布和雪山而闻名。",
        "作为世界自然遗产，九寨沟的五花海和长海是必看景点。",
        "这个位于阿坝州的景区是徒步和摄影爱好者的天堂。",
    ],
    "乐山大佛": [
        "乐山大佛是一尊巨大的石刻弥勒佛坐像，濒临岷江。",
        "游客可以乘船远观大佛全景，也可以近距离攀爬参观。",
        "这尊佛像开凿于唐代，历时约90年才最终完成。",
    ],
    "成都大熊猫基地": [
        "成都大熊猫繁育研究基地是观赏中国国宝的最佳地点。",
        "在这里，你可以看到从幼崽到成年大熊猫的各种形态。",
        "该基地致力于大熊猫的保护和繁育工作。",
    ]
}

# 创建三元组列表
triplets = []
attractions = list(data.keys())

for attraction_name, sentences in data.items():
    # 从一个景点中至少选出两个句子作为 anchor 和 positive
    if len(sentences) < 2:
        continue

    # 创建正样本对
    for i in range(len(sentences)):
        for j in range(i + 1, len(sentences)):
            anchor = sentences[i]
            positive = sentences[j]

            # 随机选择一个不同的景点作为负样本来源
            negative_attraction_name = random.choice([name for name in attractions if name != attraction_name])
            negative = random.choice(data[negative_attraction_name])

            triplets.append(InputExample(texts=[anchor, positive, negative]))

print(f"成功创建 {len(triplets)} 个三元组样本。")
# 打印一个样本看看
print("示例三元组:", triplets[0].texts)

# --- 步骤 2: 选择一个预训练模型 ---
print("\n--- 2. 选择预训练模型 ---")
model_name = 'BAAI/bge-small-zh-v1.5'
# 如果你想让模型输出固定长度的向量，可以添加一个池化层
# word_embedding_model = models.Transformer(model_name)
# pooling_model = models.Pooling(word_embedding_model.get_word_embedding_dimension())
# model = SentenceTransformer(modules=[word_embedding_model, pooling_model])
model = SentenceTransformer(model_name)
print(f"已加载模型: {model_name}")

# --- 步骤 3: 定义数据加载器和损失函数 ---
print("\n--- 3. 定义数据加载器和损失函数 ---")
# TripletLoss 需要一个 DataLoader，它每次返回一个批次的三元组
# NoDuplicatesDataLoader 确保一个批次内没有重复的句子，这对于三元组损失很重要
train_dataloader = DataLoader(triplets, shuffle=True, batch_size=4)
# 定义损失函数。使用余弦距离计算
train_loss = losses.TripletLoss(model=model)
print("已定义 TripletLoss 和 DataLoader。")

# --- 步骤 4: 微调模型 ---
print("\n--- 4. 开始微调模型 ---")
# 训练参数
epochs = 2
warmup_steps = int(len(train_dataloader) * epochs * 0.1)  # 10% 的步数用于预热

# 开始训练
model.fit(train_objectives=[(train_dataloader, train_loss)],
          epochs=epochs,
          warmup_steps=warmup_steps,
          output_path='./sichuan_attractions_model',  # 指定模型保存路径
          show_progress_bar=True)

print("--- 模型微调完成并已保存！ ---")

# --- 步骤 5: 在推荐系统中使用微调后的模型 ---
print("\n--- 5. 使用微调后的模型进行推荐 ---")


